#!/usr/bin/env python3
# -*- encoding: utf-8 -*-
"""
@File         :switch_fw_mode.py
@Description  :
@Time         :2024/05/29 22:14:36
@Author       :Lin Yuheng
@Version      :1.0
"""

import rospy
from mavros_msgs.msg import State
from mavros_msgs.srv import SetMode, CommandBool, SetModeRequest, SetModeResponse


class SwitchFWMode:
    def __init__(self):
        self.planeID = 0
        self.uavID = "plane_0"

        self.current_state = State()
        self.mode_request = SetModeRequest()

        self.outflag = False  # 退出标志位

        self.counters = 0  # 计数器
        self.times_out = 0  # 最大计数次数

    def ros_sub_pub(self):
        # Subscriber
        # 无人机当前状态
        self.state_sub = rospy.Subscriber(
            self.uavID + "/mavros/state",
            State,
            callback=self.fw_states_callback,
            queue_size=10,
        )

        # srv
        self.arm_service = rospy.ServiceProxy(
            self.uavID + "/mavros/cmd/arming",
            CommandBool,
        )

        self.set_mode_service = rospy.ServiceProxy(
            self.uavID + "/mavros/set_mode", SetMode
        )

    def set_planeID(self, id:int):
        self.planeID = id
        if self.planeID == 0:
            self.uavID = "plane_0"
        elif self.planeID == 1:
            self.uavID = "plane_1"
        elif self.planeID == 2:
            self.uavID = "plane_2"
        elif self.planeID == 3:
            self.uavID = "plane_3"

    def fw_states_callback(self, msg: State):
        self.current_state = msg

    def get_ros_time(self, begin: rospy.Time):
        time_now = rospy.Time.now()
        currTimeSec = time_now.secs - begin.secs
        currTimenSec = time_now.nsecs / 1e9 - begin.nsecs / 1e9
        return currTimeSec + currTimenSec

    def arm(self):
        if self.arm_service(True):
            return True
        else:
            return False

    def start(self):
        begin_time = rospy.Time.now()
        current_time = self.get_ros_time(begin_time)
        rate = rospy.Rate(5)
        self.ros_sub_pub()

        while not rospy.is_shutdown() and not self.outflag:
            print(
                ">>>>>>>>>>>>>>>>>>>>>>>>>>>>>>SWITCH_FW_MODE<<<<<<<<<<<<<<<<<<<<<<<<<<"
            )
            rospy.loginfo("本机ID:%d", self.planeID)
            rospy.loginfo("当前时刻:%f", current_time)
            rospy.loginfo("当前模式 : [ %s ]", self.current_state.mode)
            rospy.loginfo("切换模式: 1 arm, 2 takeoff, 3 offboard, 4 loiter, 5 out")
            # 获取键盘指令

            try:
                cmd = input("请输入指令:")

                if cmd == "1":
                    self.arm_service = self.arm()
                elif cmd == "2":
                    self.mode_request.custom_mode = "AUTO.TAKEOFF"
                elif cmd == "3":
                    self.mode_request.custom_mode = "OFFBOARD"
                elif cmd == "4":
                    self.mode_request.custom_mode = "AUTO.LOITER"
                elif cmd == "5":
                    self.outflag = True
                    break
                else:
                    print("无效指令，重新输入")

                while (
                    cmd in ['2','3','4']
                    and (self.mode_request.custom_mode != self.current_state.mode)
                    and (self.counters <= self.times_out)
                ):
                    self.counters += 1
                    self.set_mode_service.call(self.mode_request)

                self.counters = 0
                rate.sleep()

            except rospy.ROSException as e:
                pass

        rospy.loginfo("switch_fw_mode end")


if __name__ == "__main__":
    rospy.init_node("switch_fw_node")
    switch_fw_mode = SwitchFWMode()
    switch_fw_mode.start()
